{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 额外安装一个库  pip install rouge-chinese"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = '/data/datasets/supremezxc/nlpcc_data.json'\n",
    "model_dir = '/data/models/huggingface/mengzi-t5-base'\n",
    "save_dir = '/data/logs/abstract_model'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入相关包\n",
    "import torch\n",
    "from datasets import Dataset,load_dataset\n",
    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,T5ForConditionalGeneration,DataCollatorForSeq2Seq,Seq2SeqTrainer,Seq2SeqTrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载数据集\n",
    "datasets = load_dataset('json', data_files=data_dir,split='train')\n",
    "datasets = datasets.train_test_split(test_size=0.90)\n",
    "datasets = datasets[\"train\"]\n",
    "print(datasets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = datasets.train_test_split(100,seed=42)\n",
    "datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets[\"train\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据处理\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_func(examples):\n",
    "    contents = [\"内容扩充:\\n\" + e for e in examples[\"title\"]]\n",
    "    # 分开处理\n",
    "    inputs =  tokenizer(contents, max_length=128, truncation=True)\n",
    "    labels =  tokenizer(text_target=examples[\"content\"], max_length=512, truncation=True)\n",
    "    inputs[\"labels\"] =  labels[\"input_ids\"]\n",
    "    return inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_ds =  datasets.map(process_func,batched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 将数据处理成如下这个样子\n",
    "print(tokenizer.decode(tokenizer_ds[\"train\"][0][\"input_ids\"]))\n",
    "print(tokenizer.decode(tokenizer_ds[\"train\"][0][\"labels\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建模型\n",
    "model =  AutoModelForSeq2SeqLM.from_pretrained(model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建评估函数\n",
    "import numpy as np\n",
    "from rouge_chinese import Rouge\n",
    "rouge = Rouge()\n",
    "\n",
    "def compute_metric(eval_preds):\n",
    "    predictions,labels = eval_preds\n",
    "    decoded_preds = tokenizer.batch_decode(predictions,skip_special_tokens=True)\n",
    "    # labels 有-100 因此需要进行填充操作,因为labels有-100 不属于特殊字符，因此需要填充为特殊字符\n",
    "    labels = np.where(labels!=-100,labels,tokenizer.pad_token_id)\n",
    "    decoded_labels = tokenizer.batch_decode(labels,skip_special_tokens=True)\n",
    "    # 基于中文的字来做评估\n",
    "    decoded_preds = [\" \".join(pred.strip()) for pred in decoded_preds]\n",
    "    decoded_labels = [\" \".join(label.strip()) for label in decoded_labels]\n",
    "    # 该方法计算 Rouge，标准的文字评估方案\n",
    "    scores = rouge.get_scores(decoded_preds,decoded_labels,avg=True)\n",
    "    return {\n",
    "        \"rouge-1\":scores[\"rouge-1\"][\"f\"],\n",
    "        \"rouge-2\":scores[\"rouge-2\"][\"f\"],\n",
    "        \"rouge-l\":scores[\"rouge-l\"][\"f\"]\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = Seq2SeqTrainingArguments(\n",
    "    output_dir=save_dir,\n",
    "    per_device_train_batch_size=2,\n",
    "    per_device_eval_batch_size=2,\n",
    "    gradient_accumulation_steps=2,\n",
    "    logging_steps = 100,\n",
    "    num_train_epochs=4,\n",
    "    eval_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    metric_for_best_model=\"rouge-l\",\n",
    "    # predict_with_generate 这个参数必须设置为True，否则无法做评估\n",
    "    predict_with_generate = True\n",
    ")\n",
    "args.num_train_epochs              # total number of training epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer =  Seq2SeqTrainer(\n",
    "    args = args,\n",
    "    model =model,\n",
    "    train_dataset = tokenizer_ds[\"train\"],\n",
    "    eval_dataset = tokenizer_ds[\"test\"],\n",
    "    compute_metrics = compute_metric,\n",
    "    tokenizer =tokenizer,\n",
    "    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer)\n",
    "\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 开始训练\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n"
     ]
    }
   ],
   "source": [
    "# 模型推理\n",
    "# 导入相关包\n",
    "import torch\n",
    "import os\n",
    "from datasets import Dataset,load_dataset\n",
    "from transformers import pipeline\n",
    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,T5ForConditionalGeneration,DataCollatorForSeq2Seq,Seq2SeqTrainer,Seq2SeqTrainingArguments\n",
    "\n",
    "# data_dir = '/data/datasets/supremezxc/nlpcc_data.json'\n",
    "model_dir = '/data/models/huggingface/mengzi-t5-base'\n",
    "save_dir = '/data/logs/abstract_model'\n",
    "# E:\\data\\logs\\abstract_model\\checkpoint-4900\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_dir)\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(os.path.join(save_dir,'checkpoint-4900')).to('cuda')\n",
    "pipe = pipeline(\"text2text-generation\",model=model,tokenizer=tokenizer,device=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#pipe(\"内容扩充:\\n\"+datasets[\"test\"][-2][\"content\"],max_length=512,do_sample=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = '我看到一个苹果，该苹果的位置是在（100，200，333）'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'generated_text': '新浪科技讯北京时间5月1日消息,今天(9月12日),我在手机屏幕上看到苹果,该苹果在Android的右侧了,但是没有经过我指指点点,就没有点数,因为位置在这里,于是就出现了这个苹果,不知道应该是在(200,200,2.3,333)的位置呢?我首先确认是:某iPhone6,这个苹果的位置是在(1000,200,333)333;然后看到它的位置是在(30,220,600(220,600),400,333;(333)在苹果的左边位置是:3,39;(5)。在距离Cid上,iPhone5(90,800,34;33,300,400,333)338/15。(手机应用商店(电脑)搜索关注公众号“参考消息”(ID:ckxxwx),外国媒体每日报道精选,随时随地想看就看,还有会员福利等着您哦~或者关注奥一网官方微信。(微信号:oeeeend)】您若对某微信公众号有任何疑问或评论,或者关注奥一网官方微信:oeeend)或关注同微信号:oeeeend;或关注奥一网官方微信。'}]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pipe(\"内容扩充:\\n\"+text,max_length=1800,do_sample=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
